# @Version        : 1.0
# @Update Time    : 2025/3/17 21:38
# @File           : db_database.py
# @IDE            : PyCharm
# @Desc           : 数据库配置和依赖注入
from typing import AsyncGenerator, Optional, Dict, Any, Callable
from advanced_alchemy.config import EngineConfig
from advanced_alchemy.extensions.fastapi import (
    AdvancedAlchemy,
    AsyncSessionConfig,
    SQLAlchemyAsyncConfig,
)
from sqlalchemy.ext.asyncio import AsyncSession
from src.config import settings
from src.log import logger

__all__ = ["alchemy", "sqlalchemy_config", "get_depend_session", "DatabaseConfig"]


class DatabaseConfig:
    """数据库配置类"""

    def __init__(
        self,
        echo: bool = True,
        echo_pool: bool = False,
        pool_pre_ping: bool = True,
        pool_recycle: int = 3600,
        pool_size: int = 5,
        max_overflow: int = 5,
        expire_on_commit: bool = False,
        autoflush: bool = False,
    ):
        self.engine_config = EngineConfig(
            echo=echo,
            echo_pool=echo_pool,
            pool_pre_ping=pool_pre_ping,
            pool_recycle=pool_recycle,
            pool_size=pool_size,
            max_overflow=max_overflow,
        )
        self.session_config = AsyncSessionConfig(
            expire_on_commit=expire_on_commit,
            autoflush=autoflush,
        )


class DatabaseManager:
    """数据库管理器"""

    def __init__(self, config: Optional[DatabaseConfig] = None):
        self.config = config or DatabaseConfig()
        self._alchemy: Optional[AdvancedAlchemy] = None
        self._configs: Dict[str, SQLAlchemyAsyncConfig] = {}
        self._init_configs()

    def _init_configs(self) -> None:
        """初始化数据库配置"""
        try:
            for key, value in settings.db.ORM_DB.DATA_SOURCES.items():
                self._configs[key] = SQLAlchemyAsyncConfig(
                    bind_key=key,
                    connection_string=value.DB_URL,
                    session_config=self.config.session_config,
                    create_all=False,
                    engine_config=self.config.engine_config,
                )
            self._alchemy = AdvancedAlchemy(config=list(self._configs.values()))
            logger.info("数据库配置初始化成功")
        except Exception as e:
            print(f"数据库配置初始化失败: {e}")
            raise

    @property
    def alchemy(self) -> AdvancedAlchemy:
        """获取数据库实例"""
        if self._alchemy is None:
            self._init_configs()
        return self._alchemy

    def get_session_dependency(
        self, bind_key: Optional[str] = None
    ) -> Callable[[], AsyncGenerator[AsyncSession | None, Any]]:
        """
        获取数据库会话依赖
        :param bind_key: 数据库绑定键
        :return: 异步数据库会话
        """

        async def get_db_async():
            session: AsyncSession | None = None
            try:
                async with self.alchemy.with_async_session(key=bind_key) as session:
                    yield session
            except Exception as e:
                print(f"数据库会话创建失败: {e}")
                if session:
                    try:
                        await session.close()
                    except Exception as close_error:
                        print(f"关闭数据库会话失败: {close_error}")
                raise
            finally:
                if session:
                    try:
                        await session.close()
                    except Exception as close_error:
                        print(f"关闭数据库会话失败: {close_error}")

        return get_db_async


# 创建默认配置实例
default_config = DatabaseConfig(echo=False)
db_manager = DatabaseManager(config=default_config)

# 导出常用实例和函数
alchemy = db_manager.alchemy
sqlalchemy_config = (
    list(db_manager._configs.values())[0] if db_manager._configs else None
)
get_depend_session = db_manager.get_session_dependency
